import asyncio
import aiohttp
from tqdm import tqdm
import os
import sys
from typing import List, Optional, Dict, Tuple
from dataclasses import dataclass, field
from asyncio import Semaphore
import math
from datetime import datetime
import shutil
import signal


@dataclass
class ProxyConfig:
    """代理配置类"""
    url: str
    auth: Optional[Dict[str, str]] = None
    headers: Optional[Dict[str, str]] = None
    
    def to_proxy_settings(self) -> Dict:
        settings = {'proxy': self.url}
        if self.auth:
            settings['proxy_auth'] = aiohttp.BasicAuth(**self.auth)
        if self.headers:
            settings['proxy_headers'] = self.headers
        return settings

@dataclass
class DownloadChunk:
    """文件分块下载的块信息"""
    start: int
    end: int
    downloaded: int = 0
    completed: bool = False

@dataclass
class DownloadTask:
    url: str
    directory: str
    filename: str
    proxy_config: Optional[ProxyConfig] = None
    total_size: int = 0
    chunks: List[DownloadChunk] = None
    progress_bar: tqdm = None
    
    def get_filename(self) -> str:
        """获取下载文件名"""
        return f"{self.directory}/{self.filename}"  

    def get_directory(self) -> str:
        """获取下载目录"""
        return self.directory
    
    def get_temp_filename(self) -> str:
        """获取临时文件名"""
        return f"{self.directory}/{self.filename}.downloading"
    
    def get_chunk_filename(self, chunk_id: int) -> str:
        """获取分块临时文件名"""
        return f"{self.directory}/{self.filename}.part{chunk_id}"

@dataclass
class DownloadStats:
    """下载统计信息"""
    total_files: int = 0
    total_size: int = 0
    completed_files: int = 0
    failed_files: int = 0
    skipped_files: int = 0  # 新增：跳过下载的文件数
    start_time: datetime = None
    end_time: datetime = None
    failed_files_info: Dict[str, str] = field(default_factory=dict)
    directories: set = field(default_factory=set)

class DownloadManager:
    def __init__(self, 
                 max_concurrent_files: int = 3,
                 max_concurrent_chunks: int = 5,
                 chunk_size: int = 100*1024*1024,
                 default_proxy: Optional[ProxyConfig] = None, 
                 hf_token: str = None,
                 max_retries: int = 20,          # 最大重试次数
                 retry_delay: int = 20,         # 重试等待时间（秒）
                 timeout: int = 600):           # 连接超时时间（秒）
        """
        初始化下载管理器
        :param max_concurrent_files: 同时下载的最大文件数(1-10)
        :param max_concurrent_chunks: 单个文件的最大分块数(1-5)
        :param chunk_size: 触发分块下载的文件大小阈值
        :param max_retries: 下载失败时的最大重试次数
        :param retry_delay: 重试之间的等待时间
        :param timeout: 连接超时时间
        """
        # 限制并发文件数在1-10之间
        if max_concurrent_files < 1:
            print(f"警告：max_concurrent_files({max_concurrent_files})小于1，已设置为1")
            max_concurrent_files = 1
        elif max_concurrent_files > 10:
            print(f"警告：max_concurrent_files({max_concurrent_files})大于10，已设置为10")
            max_concurrent_files = 10
        
        # 限制分块数在1-5之间
        if max_concurrent_chunks < 1:
            print(f"警告：max_concurrent_chunks({max_concurrent_chunks})小于1，已设置为1")
            max_concurrent_chunks = 1
        elif max_concurrent_chunks > 5:
            print(f"警告：max_concurrent_chunks({max_concurrent_chunks})大于5，已设置为5")
            max_concurrent_chunks = 5
        
        self.semaphore = Semaphore(max_concurrent_files)
        self.tasks: List[DownloadTask] = []
        self.default_proxy = default_proxy
        self.hf_token = hf_token
        self.chunk_size = chunk_size
        self.max_concurrent_chunks = max_concurrent_chunks
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        self.timeout = aiohttp.ClientTimeout(
            total=timeout,
            connect=60,
            sock_connect=60,
            sock_read=timeout
        )
        self.stats = DownloadStats()
        self.overall_progress: Optional[tqdm] = None        # 总进度条
        self.total_downloaded_size: int = 0                 # 已下载的总大小
        self._running = True
        self._tasks = []
    
    async def add_task(self, url: str, directory: str, filename: str, total_size : int = 0, proxy_config: Optional[ProxyConfig] = None):
        """添加下载任务，可以为单个任务指定代理配置"""
        task = DownloadTask(url=url, directory=directory, filename=filename, total_size=total_size, proxy_config=proxy_config or self.default_proxy)
        self.tasks.append(task)
        self.stats.directories.add(directory)
        self.stats.total_files += 1
        
    async def get_file_info_with_retry(self, session: aiohttp.ClientSession, url: str, proxy_settings: dict) -> Tuple[int, str]:
        """带重试机制的获取文件信息"""
        last_error = None
        for retry in range(self.max_retries):
            try:
                headers = {
                    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36',
                    'Accept': '*/*',
                    'Accept-Encoding': 'gzip, deflate, br',
                    'Connection': 'keep-alive'
                }
                if self.hf_token:
                    headers['Authorization'] = f'Bearer {self.hf_token}'

                # 直接使用GET请求而不是HEAD请求
                async with session.get(url, headers=headers, **proxy_settings, ssl=False) as response:
                    if response.status in [200, 206]:
                        # 优先使用content-length
                        size = int(response.headers.get('content-length', 0))
                        if size == 0:
                            # 如果content-length不存在，读取整个响应内容来获取大小
                            content = await response.read()
                            size = len(content)
                        etag = response.headers.get('etag', '')
                        return size, etag
            
            except Exception as e:
                last_error = e
                if retry < self.max_retries - 1:
                    wait_time = (retry + 1) * self.retry_delay
                    print(f"获取文件信息失败: {str(e)}, 等待 {wait_time} 秒后重试... ({retry + 1}/{self.max_retries})")
                    await asyncio.sleep(wait_time)
                else:
                    print(f"获取文件信息失败: {str(e)}, 已达到最大重试次数")
        return 0, ''

    async def check_file_complete(self, filename: str, expected_size: int) -> bool:
        """检查文件是否已完整下载"""
        if not os.path.exists(filename):
            return False
        
        actual_size = os.path.getsize(filename)
        if actual_size == expected_size:
            # 更新总进度条（跳过的文件也计入进度）
            if self.overall_progress:
                self.overall_progress.update(1)
                self.overall_progress.set_description(
                    f">>> 总进度 [{self.stats.completed_files + self.stats.skipped_files + 1}/{self.stats.total_files}]"
                )
                # self.overall_progress.set_postfix({" ":" "*20})
            return True
        return False

    async def download_with_retry(self, session: aiohttp.ClientSession, url: str, 
                                headers: dict, proxy_settings: dict) -> aiohttp.ClientResponse:
        """带重试机制的下载请求"""
        last_error = None
        for retry in range(self.max_retries):
            try:
                response = await session.get(url, headers=headers, **proxy_settings, ssl=False)
                
                # 修改：200和206都是有效的响应
                if response.status in [200, 206]:
                    return response
                
                # 处理其他状态码
                if response.status == 429:  # Too Many Requests
                    wait_time = int(response.headers.get('Retry-After', self.retry_delay))
                    print(f"请求频率限制，等待 {wait_time} 秒后重试...")
                    await asyncio.sleep(wait_time)
                elif response.status >= 500:  # 服务器错误，可以重试
                    print(f"服务器错误，状态码: {response.status}，重试中... ({retry + 1}/{self.max_retries})")
                    await asyncio.sleep(self.retry_delay)
                else:  # 其他错误状态码
                    raise aiohttp.ClientError(f"请求失败，状态码: {response.status}")
                
                # 确保响应被正确关闭
                if not response.closed:
                    await response.release()
                
            except (aiohttp.ClientError, asyncio.TimeoutError) as e:
                last_error = e
                if retry < self.max_retries - 1:
                    wait_time = (retry + 1) * self.retry_delay
                    print(f"\n连接错误: {str(e)}，等待 {wait_time} 秒后重试... ({retry + 1}/{self.max_retries})\n")
                    await asyncio.sleep(wait_time)
                else:
                    raise Exception(f"重试{self.max_retries}次后仍然失败: {str(e)}")
                
        raise Exception(f"达到最大重试次数: {str(last_error)}")

    async def download_chunk_with_retry(self, session: aiohttp.ClientSession, task: DownloadTask, 
                                      chunk: DownloadChunk, chunk_id: int, proxy_settings: dict):
        """带重试机制的分块下载"""
        chunk_file = task.get_chunk_filename(chunk_id)
        expected_size = chunk.end - chunk.start + 1
        
        # 检查是否存在下载目录
        if not os.path.exists(task.get_directory()):
            os.makedirs(task.get_directory())
        
        # 检查已下载的部分
        if os.path.exists(chunk_file):
            chunk.downloaded = os.path.getsize(chunk_file)
            if chunk.downloaded == expected_size:
                chunk.completed = True
                task.progress_bar.update(expected_size - task.progress_bar.n)
                return
            else:
                os.remove(chunk_file)
                chunk.downloaded = 0

        headers = {
            'Range': f'bytes={chunk.start}-{chunk.end}',
            'Connection': 'keep-alive'
        }
        
        for retry in range(self.max_retries):
            try:
                async with await self.download_with_retry(session, task.url, headers, proxy_settings) as response:
                    # 修改：处理服务器不支持Range请求的情况
                    if response.status == 200:
                        # 服务器返回完整文件，需要手动截取所需部分
                        temp_chunk_file = f"{chunk_file}.tmp"
                        try:
                            total_read = 0
                            with open(temp_chunk_file, 'wb') as f:
                                async for data in response.content.iter_chunked(8192):
                                    # 跳过不需要的部分
                                    if total_read < chunk.start:
                                        if total_read + len(data) <= chunk.start:
                                            total_read += len(data)
                                            continue
                                        else:
                                            # 部分数据需要写入
                                            offset = chunk.start - total_read
                                            data = data[offset:]
                                            total_read += offset
                                    
                                    # 写入需要的部分
                                    if total_read + len(data) > chunk.end + 1:
                                        # 截取到需要的长度
                                        data = data[:chunk.end - total_read + 1]
                                    
                                    f.write(data)
                                    data_len = len(data)
                                    chunk.downloaded += data_len
                                    task.progress_bar.update(data_len)
                                    total_read += data_len
                                    
                                    if total_read >= chunk.end + 1:
                                        break
                                    
                                task.progress_bar.set_description(
                                    f"下载 {task.filename} [块 {chunk_id + 1}/{len(task.chunks)}]"
                                )
                            
                            # 验证下载的分块大小
                            actual_size = os.path.getsize(temp_chunk_file)
                            if actual_size != expected_size:
                                raise ValueError(f"分块{chunk_id}下载完成后大小不匹配：预期{expected_size}字节，实际{actual_size}字节")
                            
                            os.replace(temp_chunk_file, chunk_file)
                            chunk.completed = True
                            return
                            
                        except Exception as e:
                            if os.path.exists(temp_chunk_file):
                                os.remove(temp_chunk_file)
                            raise
                    
                    elif response.status == 206:
                        # 原有的处理逻辑保持不变
                        content_range = response.headers.get('Content-Range', '')
                        if not content_range:
                            raise ValueError("服务器响应中没有Content-Range头")
                        
                        # 解析Content-Range: bytes start-end/total
                        try:
                            range_info = content_range.split(' ')[1].split('/')[0]
                            range_start, range_end = map(int, range_info.split('-'))
                            if range_start != chunk.start or range_end != chunk.end:
                                raise ValueError(f"服务器返回的范围({range_start}-{range_end})与请求的范围({chunk.start}-{chunk.end})不匹配")
                        except Exception as e:
                            raise ValueError(f"无法解析Content-Range: {content_range}, 错误: {str(e)}")
                        
                        # 使用临时文件下载
                        temp_chunk_file = f"{chunk_file}.tmp"
                        try:
                            with open(temp_chunk_file, 'wb') as f:
                                async for data in response.content.iter_chunked(8192):
                                    f.write(data)
                                    data_len = len(data)
                                    chunk.downloaded += data_len
                                    task.progress_bar.update(data_len)
                            task.progress_bar.set_description(
                                f"下载 {task.filename} [块 {chunk_id + 1}/{len(task.chunks)}]"
                            )
                        
                            # 验证下载的分块大小
                            actual_size = os.path.getsize(temp_chunk_file)
                            if actual_size != expected_size:
                                raise ValueError(f"分块{chunk_id}下载完成后大小不匹配：预期{expected_size}字节，实际{actual_size}字节")
                            
                            # 下载完成后重命名临时文件
                            os.replace(temp_chunk_file, chunk_file)
                            chunk.completed = True
                            return
                        except Exception as e:
                            if os.path.exists(temp_chunk_file):
                                os.remove(temp_chunk_file)
                            raise
                    
                    else:
                        raise ValueError(f"服务器返回了意外的状态码: {response.status}")
                
            except Exception as e:
                if retry < self.max_retries - 1:
                    wait_time = (retry + 1) * self.retry_delay
                    print(f"{task.filename}: 分块 {chunk_id} 下载失败: {str(e)}, 等待 {wait_time} 秒后重试... ({retry + 1}/{self.max_retries})")
                    await asyncio.sleep(wait_time)
                else:
                    print(f"{task.filename}: 分块 {chunk_id} 下载失败: {str(e)}, 已达到最大重试次数")
                    raise

    async def safe_rename(self, src: str, dst: str) -> bool:
        """安全地重命名文件，如果目标文件存在则先删除"""
        try:
            if os.path.exists(dst):
                os.remove(dst)
            os.rename(src, dst)
            return True
        except Exception as e:
            print(f"重命名文件失败 {src} -> {dst}, 错误: {str(e)}")
            return False

    async def download_with_chunks(self, session: aiohttp.ClientSession, task: DownloadTask, proxy_settings: dict):
        """使用分块方式下载"""
        # 根据最大并发分块数计算每个分块的大小
        base_chunk_size = task.total_size // self.max_concurrent_chunks
        remainder = task.total_size % self.max_concurrent_chunks
        task.chunks = []
      
        # 创建分块，确保所有分块大小加起来等于总大小
        current_position = 0
        for i in range(self.max_concurrent_chunks):
            # 最后一个分块处理余数
            chunk_size = base_chunk_size + (1 if i < remainder else 0)
            start = current_position
            end = start + chunk_size - 1
            current_position = end + 1
            
            # 检查分块是否已下载
            chunk_file = task.get_chunk_filename(i)
            downloaded = 0
            if os.path.exists(chunk_file):
                downloaded = os.path.getsize(chunk_file)
            
            task.chunks.append(DownloadChunk(start, end, downloaded))
        
        # 验证分块总大小
        total_chunks_size = sum(chunk.end - chunk.start + 1 for chunk in task.chunks)
        if total_chunks_size != task.total_size:
            raise ValueError(f"分块总大小({total_chunks_size})与文件大小({task.total_size})不匹配")
        
        # 并发下载所有分块
        while True:
            incomplete_chunks = [(i, chunk) for i, chunk in enumerate(task.chunks) if not chunk.completed]
            if not incomplete_chunks:
                break
                
            chunk_tasks = []
            for i, chunk in incomplete_chunks:
                chunk_tasks.append(self.download_chunk_with_retry(session, task, chunk, i, proxy_settings))
            
            try:
                await asyncio.gather(*chunk_tasks)
            except Exception as e:
                print(f"部分分块下载失败，将在下次启动时继续下载")
                raise
            
        # 所有分块下载完成，进行合并
        try:
            await self.merge_chunks(task, self.max_concurrent_chunks)
            # 合并完成后重命名
            temp_file = task.get_temp_filename()
            if os.path.exists(temp_file):
                await self.safe_rename(temp_file, task.get_filename())
                # 清理分块文件
                for i in range(self.max_concurrent_chunks):
                    chunk_file = task.get_chunk_filename(i)
                    if os.path.exists(chunk_file):
                        os.remove(chunk_file)
            # 更新进度为100%
            task.progress_bar.update(task.total_size - task.progress_bar.n)
        except Exception as e:
            print(f"合并分块文件失败: {str(e)}")
            raise

    async def download_without_chunks(self, session: aiohttp.ClientSession, task: DownloadTask, proxy_settings: dict):
        """使用普通方式下载"""
        temp_file = task.get_temp_filename()  # .downloading
        temp_temp_file = f"{temp_file}.tmp"   # .downloading.tmp
        initial_size = 0
        
        # 检查是否存在下载目录
        if not os.path.exists(task.get_directory()):
            os.makedirs(task.get_directory())
        
        # 修改：首先检查.downloading.tmp文件
        if os.path.exists(temp_temp_file):
            initial_size = os.path.getsize(temp_temp_file)
            if initial_size >= task.total_size:
                # 如果临时文件已经完整，直接重命名为最终文件
                await self.safe_rename(temp_temp_file, task.get_filename())
                if task.progress_bar:
                    task.progress_bar.n = task.total_size
                    task.progress_bar.refresh()
                return
            print(f"\n发现临时文件，继续从 {initial_size} 字节处下载 {task.filename}")
        # 如果.tmp文件不存在，检查.downloading文件
        elif os.path.exists(temp_file):
            temp_size = os.path.getsize(temp_file)
            if temp_size >= task.total_size:
                # 如果.downloading文件已经完整，直接重命名为最终文件
                await self.safe_rename(temp_file, task.get_filename())
                if task.progress_bar:
                    task.progress_bar.n = task.total_size
                    task.progress_bar.refresh()
                return
            # 如果.downloading文件存在但不完整，复制到.tmp文件继续下载
            shutil.copy2(temp_file, temp_temp_file)
            initial_size = temp_size
            print(f"\n从.downloading文件继续下载，起始位置: {initial_size} 字节")
        
        if initial_size > 0:
            task.progress_bar.update(initial_size)
        
        for retry in range(self.max_retries):
            try:
                headers = {'Connection': 'keep-alive'}
                if initial_size > 0:
                    headers['Range'] = f'bytes={initial_size}-'
                
                async with await self.download_with_retry(session, task.url, headers, proxy_settings) as response:
                    # 处理断点续传的情况
                    if initial_size > 0 and response.status != 206:
                        print(f"\n警告：服务器不支持断点续传，将重新下载 {task.filename}")
                        initial_size = 0
                        if os.path.exists(temp_temp_file):
                            os.remove(temp_temp_file)
                    
                    mode = 'ab' if initial_size > 0 else 'wb'
                    with open(temp_temp_file, mode) as f:
                        downloaded = initial_size
                        async for chunk in response.content.iter_chunked(8192*1024):
                            if not chunk:
                                continue
                            
                            f.write(chunk)
                            chunk_size = len(chunk)
                            downloaded += chunk_size
                            task.progress_bar.update(chunk_size)
                            task.progress_bar.set_description(f"下载 {task.filename}")
                            
                            # 定期刷新文件缓冲区
                            if downloaded % (8192*1024) == 0:
                                f.flush()
                                os.fsync(f.fileno())
                    
                    # 验证下载的大小
                    actual_size = os.path.getsize(temp_temp_file)
                    if actual_size < task.total_size:
                        raise ValueError(f"下载不完整：预期{task.total_size}字节，实际{actual_size}字节")
                    elif actual_size > task.total_size:
                        print(f"\n警告：下载的文件大小({actual_size})超过预期({task.total_size})，将重试下载")
                        if os.path.exists(temp_temp_file):
                            os.remove(temp_temp_file)
                        raise ValueError("文件大小超过预期")
                    
                    # 下载成功，先重命名为.downloading，再重命名为最终文件
                    os.replace(temp_temp_file, temp_file)
                    await self.safe_rename(temp_file, task.get_filename())
                    
                    # 确保进度条显示100%
                    if task.progress_bar:
                        task.progress_bar.n = task.total_size
                        task.progress_bar.refresh()
                    return

            except aiohttp.ClientPayloadError as e:
                if retry < self.max_retries - 1:
                    wait_time = (retry + 1) * self.retry_delay
                    print(f"\n下载中断(数据传输错误): {task.filename}")
                    print(f"等待 {wait_time} 秒后重试... ({retry + 1}/{self.max_retries})")
                    await asyncio.sleep(wait_time)
                else:
                    raise
            except Exception as e:
                if retry < self.max_retries - 1:
                    wait_time = (retry + 1) * self.retry_delay
                    print(f"\n下载中断: {task.filename}, 错误: {str(e)}")
                    print(f"等待 {wait_time} 秒后重试... ({retry + 1}/{self.max_retries})")
                    await asyncio.sleep(wait_time)
                else:
                    raise

    async def merge_chunks(self, task: DownloadTask, num_chunks: int):
        """合并下载的分块"""
        temp_file = task.get_temp_filename()
        final_temp_file = f"{temp_file}.final"
        
        try:
            with open(final_temp_file, 'wb') as outfile:
                # 确保按照分块的起始位置顺序合并
                sorted_chunks = sorted(enumerate(task.chunks), key=lambda x: x[1].start)
                
                # 验证分块的连续性
                for i in range(len(sorted_chunks)-1):
                    current_chunk = sorted_chunks[i][1]
                    next_chunk = sorted_chunks[i+1][1]
                    if current_chunk.end + 1 != next_chunk.start:
                        raise ValueError(f"分块不连续: 块{i}结束于{current_chunk.end}，块{i+1}开始于{next_chunk.start}")
                
                # 合并分块
                for chunk_id, chunk in sorted_chunks:
                    chunk_file = task.get_chunk_filename(chunk_id)
                    expected_size = chunk.end - chunk.start + 1
                    actual_size = os.path.getsize(chunk_file)
                    
                    if actual_size != expected_size:
                        raise ValueError(f"分块 {chunk_id} 大小不匹配: 期望 {expected_size} 字节, 实际 {actual_size} 字节")
                    #print(f"\n合并分块 {chunk_id}: 起始位置 {chunk.start}, 结束位置 {chunk.end}, 大小 {actual_size}")
                    # 使用大块读取提高效率
                    with open(chunk_file, 'rb') as infile:
                        while True:
                            data = infile.read(8192*1024)  # 8MB buffer
                            if not data:
                                break
                            outfile.write(data)
                            outfile.flush()  # 确保数据写入磁盘
                
                    # 删除已合并的分块
                    try:
                        os.remove(chunk_file)
                    except Exception as e:
                        print(f"删除分块文件失败 {chunk_file}: {str(e)}")
            
            # 验证最终文件大小
            final_size = os.path.getsize(final_temp_file)
            if final_size != task.total_size:
                raise ValueError(f"合并后的文件大小({final_size})与预期大小({task.total_size})不匹配")
            
            # 最后重命名为临时文件
            os.replace(final_temp_file, temp_file)
            
        except Exception as e:
            if os.path.exists(final_temp_file):
                os.remove(final_temp_file)
            print(f"合并分块文件失败: {str(e)}")
            raise

    async def print_download_summary(self):
        """打印下载任务汇总信息"""
        print("\n" + "="*50)
        print("下载任务汇总:")
        print(f"待下载文件数: {self.stats.total_files}")
        # print(f"总下载大小: {humanize.naturalsize(self.stats.total_size, binary=True)}")
        print("保存目录:")
        for directory in sorted(self.stats.directories):
            print(f"  - {directory}")
        files = []
        for task in self.tasks:
            parts = task.directory.split('/')
            if len(parts) > 2:
                directory = '/'.join(parts[2:])  # 取出相对路径
                files.append(f"{directory}/{task.filename}")
            else:
                files.append(task.filename)
        print(f"\n文件列表：\n{files}")
        print("="*50 + "\n")

    async def print_download_results(self):
        """打印下载结果统计"""
        self.stats.end_time = datetime.now()
        duration = self.stats.end_time - self.stats.start_time
        hours, remainder = divmod(duration.total_seconds(), 3600)
        minutes, seconds = divmod(remainder, 60)
        
        print("\n" + "="*50)
        print("下载任务完成统计:")
        print(f"开始时间: {self.stats.start_time.strftime('%Y-%m-%d %H:%M:%S')}")
        print(f"结束时间: {self.stats.end_time.strftime('%Y-%m-%d %H:%M:%S')}")
        print(f"总耗时: {int(hours)}小时{int(minutes)}分{int(seconds)}秒")
        print(f"\n计划下载: {self.stats.total_files} 个文件")
        print(f"成功下载: {self.stats.completed_files} 个文件")
        print(f"跳过下载: {self.stats.skipped_files} 个文件")  # 显示跳过的文件数
        print(f"下载失败: {self.stats.failed_files} 个文件")
        
        if self.stats.failed_files > 0:
            print("\n失败的文件:")
            for filename, error in self.stats.failed_files_info.items():
                print(f"  - {filename}: {error}")
        print("="*50)

    async def download_file(self, task: DownloadTask):
        """下载单个文件的主函数"""
        async with self.semaphore:
            conn = aiohttp.TCPConnector(
                ssl=False,
                force_close=True,
                enable_cleanup_closed=True,
                ttl_dns_cache=300,
                limit_per_host=5
            )            
            headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36',
                'Accept': '*/*',
                'Accept-Encoding': 'gzip, deflate, br',
                'Connection': 'keep-alive'
            }
            if self.hf_token and ('huggingface.co' in task.url or 'hf-mirror.com' in task.url):
                headers['Authorization'] = f'Bearer {self.hf_token}'

            try:
                async with aiohttp.ClientSession(
                    connector=conn,
                    timeout=self.timeout,
                    headers=headers,
                    trust_env=True,
                    raise_for_status=False
                ) as session:
                    proxy_settings = task.proxy_config.to_proxy_settings() if task.proxy_config else {}
                    
                    # 获取远程文件信息（如果之前没有获取到）
                    if task.total_size <= 0:
                        task.total_size, _ = await self.get_file_info_with_retry(session, task.url, proxy_settings)
                    
                    # 更新总文件大小统计
                    self.stats.total_size += task.total_size
                    
                    # 检查文件是否已完整下载
                    if await self.check_file_complete(task.get_filename(), task.total_size):
                        print(f"\n文件 {task.filename} 已存在且完整，跳过下载")
                        self.stats.skipped_files += 1
                        self.total_downloaded_size += task.total_size  # 更新已下载大小
                        if self.overall_progress:
                            self.overall_progress.set_description(
                                f">>> 总进度 [{self.stats.completed_files + self.stats.skipped_files}/{self.stats.total_files}] " 
                            )
                            # self.overall_progress.set_postfix({" ":" "*20})
                        return
                    
                    # 修改：根据文件大小判断是否需要分块下载
                    use_chunks = task.total_size > self.chunk_size
                    
                    # 设置进度条
                    task.progress_bar = tqdm(
                        total=task.total_size,
                        initial=0,
                        unit='B',
                        unit_scale=True,
                        unit_divisor=1024,
                        desc=f"下载 {task.filename}",
                        bar_format='{desc} [{percentage:6.2f}%] |{bar:30}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]'
                    ) 
                    try:
                        if use_chunks:
                            await self.download_with_chunks(session, task, proxy_settings)
                        else:
                            await self.download_without_chunks(session, task, proxy_settings)
                    except Exception as e:
                        print(f"下载中断: {task.filename}, 错误: {str(e)}")
                        print(f"已下载部分将被保留，下次启动时将继续下载")
                        raise
                    finally:
                        if task.progress_bar:
                            task.progress_bar.close()
                    
                    # 下载成功
                    self.stats.completed_files += 1
                    self.total_downloaded_size += task.total_size  # 更新已下载大小
                    # 更新总进度条
                    if self.overall_progress:
                        self.overall_progress.update(1)
                        self.overall_progress.set_description(
                            f">>> 总进度 [{self.stats.completed_files + self.stats.skipped_files}/{self.stats.total_files}]" 
                        )
                        #self.overall_progress.set_postfix({" ":" "*20})
                    
            except Exception as e:
                # 记录失败信息
                self.stats.failed_files += 1
                self.stats.failed_files_info[task.filename] = str(e)
                print(f"下载失败: {task.url}, 错误: {str(e)}")
                # 更新总进度条（失败也计入进度）
                if self.overall_progress:
                    self.overall_progress.update(1)
                    self.overall_progress.set_description(
                        f">>> 总进度 [{self.stats.completed_files + self.stats.skipped_files}/{self.stats.total_files}]" 
                    )
                    # self.overall_progress.set_postfix({" ":" "*20})
            finally:
                await conn.close()

    def format_size(self, size: int) -> str:
        """格式化文件大小显示"""
        for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
            if size < 1024:
                return f"{size:.2f}{unit}"
            size /= 1024
        return f"{size:.2f}PB"

    async def cancel_all_tasks(self):
        """取消所有正在运行的下载任务"""
        self._running = False
        print("\n正在取消所有下载任务...")
        
        # 取消所有运行中的任务
        for task in self._tasks:
            if not task.done():
                task.cancel()
        
        # 等待所有任务完成取消
        try:
            await asyncio.gather(*self._tasks, return_exceptions=True)
        except asyncio.CancelledError:
            pass
        
        # 关闭进度条
        if self.overall_progress:
            self.overall_progress.close()
        
        # 关闭所有文件的进度条
        for task in self.tasks:
            if task.progress_bar:
                task.progress_bar.close()
        
        print("所有下载任务已取消")

    async def start_all(self):
        """启动所有下载任务"""
        self.stats.start_time = datetime.now()
        self._running = True
        
        # 打印下载任务汇总
        await self.print_download_summary()
        
        # 创建总进度条，放在最上面
        self.overall_progress = tqdm(
            total=self.stats.total_files,
            initial=0,
            desc=f">>> 总进度 [0/{self.stats.total_files}] 已下载: 0B" ,
            position=0,  # 放在最上面
            bar_format='{desc} |{bar:30}| {percentage:3.0f}%'
        )
        
        try:
            # 创建下载任务
            self._tasks = [asyncio.create_task(self.download_file(task)) for task in self.tasks]
            # 等待所有任务完成
            await asyncio.gather(*self._tasks)
        except asyncio.CancelledError:
            await self.cancel_all_tasks()
        finally:
            if self.overall_progress:
                self.overall_progress.close()
        
        if self._running:
            await self.print_download_results()


# 测试代码
async def main():
    # Proxy
    proxy = ProxyConfig(
        url="http://127.0.0.1:7890",
        headers={
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
        }
    )
    
    # 配置下载管理器
    manager = DownloadManager(
        max_concurrent_files=3,
        max_concurrent_chunks=3,
        chunk_size=50*1024*1024,   # 50MB
        default_proxy=None,
        hf_token="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",
        timeout=600,               # 增加超时时间到10分钟
        max_retries=30,           # 增加重试次数
        retry_delay=20            # 增加重试等待时间
    )
       

    # 添加下载任务
    test_files = [
        ('https://www.modelscope.cn/datasets/open-thoughts/OpenThoughts-114k/resolve/master/.gitattributes', 'open-thoughts/OpenThoughts-114k', '.gitattributes', 2461),
        ('https://www.modelscope.cn/datasets/open-thoughts/OpenThoughts-114k/resolve/master/data/train-00000-of-00006.parquet', 'open-thoughts/OpenThoughts-114k/data', 'train-00000-of-00006.parquet', 250112845),
        # ('https://www.modelscope.cn/datasets/open-thoughts/OpenThoughts-114k/resolve/master/data/train-00001-of-00006.parquet', 'open-thoughts/OpenThoughts-114k/data', 'train-00001-of-00006.parquet', 174927578),
    ]
    
    for url, directory, filename, total_size in test_files:
        await manager.add_task(url, directory, filename, total_size)
    
    # 开始下载
    await manager.start_all()


if __name__ == "__main__":
    asyncio.run(main()) 